import collections


class Solution:
    def numberOfGoodPaths(self, vals, edges) -> int:
        n = len(vals)
        st = [[vals[i], i] for i in range(n)]
        st.sort()
        p = [i for i in range(n)]
        rank = [1 for _ in range(n)]
        self.ans = n
        d_edges = collections.defaultdict(set)
        for u, v in edges:
            d_edges[u].add(v)
            d_edges[v].add(u)

        def find(x):
            if p[x] == x:
                return p[x]
            else:
                p[x] = find(p[x])
                return p[x]

        def union(x, y):
            ux = find(x)
            uy = find(y)
            if ux != uy:
                if vals[ux] == vals[uy]:
                    self.ans += rank[ux] * rank[uy]
                if vals[ux] == vals[uy]:
                    p[uy] = ux
                    rank[ux] += rank[uy]
                else:
                    if vals[ux] > vals[uy]:
                        p[uy] = ux
                    else:
                        p[ux] = uy

        for num, u in st:
            for v in d_edges[u]:
                if num >= vals[v]:
                    union(u, v)

        return self.ans


data = Solution()
vals = [1, 3, 2, 1, 3]
edges = [[0, 1], [0, 2], [2, 3], [2, 4]]
print(data.numberOfGoodPaths(vals, edges))
vals = [1, 1, 2, 2, 3]
edges = [[0, 1], [1, 2], [2, 3], [2, 4]]
print(data.numberOfGoodPaths(vals, edges))
vals = [2,2,5,5]
edges = [[1,0],[0,2],[3,2]]
print(data.numberOfGoodPaths(vals, edges))
vals = [2,5,5,1,5,2,3,5,1,5]
edges = [[0,1],[2,1],[3,2],[3,4],[3,5],[5,6],[1,7],[8,4],[9,7]]
print(data.numberOfGoodPaths(vals, edges))


